{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# One-Dimensional Hopper"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook Setup \n",
    "The following cell will install Drake, checkout the underactuated repository, and set up the path (only if necessary).\n",
    "- On Google's Colaboratory, this **will take approximately two minutes** on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  Colab will ask you to \"Reset all runtimes\"; say no to save yourself the reinstall.\n",
    "- On Binder, the machines should already be provisioned by the time you can run this; it should return (almost) instantly.\n",
    "\n",
    "More details are available [here](http://underactuated.mit.edu/drake.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    import pydrake\n",
    "    import underactuated\n",
    "except ImportError:\n",
    "    !curl -s https://raw.githubusercontent.com/RussTedrake/underactuated/master/scripts/setup/jupyter_setup.py > jupyter_setup.py\n",
    "    from jupyter_setup import setup_underactuated\n",
    "    setup_underactuated()\n",
    "\n",
    "# Setup matplotlib.\n",
    "from IPython import get_ipython\n",
    "if get_ipython() is not None: get_ipython().run_line_magic(\"matplotlib\", \"inline\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# python libraries\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import HTML\n",
    "\n",
    "# pydrake imports\n",
    "from pydrake.all import (MultibodyPlant, AddMultibodyPlantSceneGraph,\n",
    "                         DiagramBuilder, Parser, VectorSystem, SignalLogger,\n",
    "                         Simulator, PlanarSceneGraphVisualizer, Multiplexer,\n",
    "                         plot_system_graphviz, MatrixGain)\n",
    "\n",
    "# underactuated imports\n",
    "from underactuated import FindResource"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem Description\n",
    "\n",
    "In this notebook we synthesize a controller for the one dimensional hopper depicted in the figure below.\n",
    "This is a simplified model of the hopper analyzed in Marc Raibert's paper \"[Hopping in legged systems - Modeling and simulation for the two-dimensional one-legged case](https://ieeexplore.ieee.org/abstract/document/6313238)\".\n",
    "Even though we reported below all the mathematical derivations you need to complete this exercise, you're encouraged to read the paper!\n",
    "\n",
    "<img src=\"https://raw.githubusercontent.com/RussTedrake/underactuated/master/figures/exercises/one_d_hopper.jpg\" width=\"500\">\n",
    "\n",
    "The hopper is composed of two rigid links: the **body** and the **leg**, with mass $m_{\\mathrm{b}}$ and $m_{\\mathrm{l}}$, respectively.\n",
    "These are connected by an elastic prismatic joint.\n",
    "The joint spring has stiffness $k$ and can be compressed in two ways: by moving the hopper foot towards the body (i.e. increasing $l (t) = z_{\\mathrm{l}} (t) - z_{\\mathrm{b}} (t)$), or by increasing the preload $\\delta (t)$, which is our (nonnegative) control variable.\n",
    "The total compression of the spring is $l (t) + \\delta (t)$, and the figure in the left below depicts the spring at rest.\n",
    "\n",
    "The control task is to make the system hop at a given height $h$.\n",
    "We will accomplish this via a careful analysis of the system energy in the various phases of the hopping cycle.\n",
    "\n",
    "Besides understanding the physics of this system, in this notebook you are asked to complete two pieces of code:\n",
    "- A function that computes the hopper energy given its state.\n",
    "- The hopping controller which sets the preload $\\delta (t)$ based on the measurement of the hopper state.\n",
    "\n",
    "In the following cell we wrote a convenient function to access the physical parameters of the hopper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_hopper_parameters():\n",
    "    \n",
    "    # spring stiffness (not in the urdf)\n",
    "    k = 1000\n",
    "    \n",
    "    # parse urdf\n",
    "    hopper = MultibodyPlant(time_step=0)\n",
    "    urdf_path = FindResource('models/one_d_hopper.urdf')\n",
    "    Parser(hopper).AddModelFromFile(urdf_path)\n",
    "    hopper.Finalize()\n",
    "    \n",
    "    # retrieve physical parameters\n",
    "    mb = hopper.GetBodyByName('body').default_mass()\n",
    "    ml = hopper.GetBodyByName('leg').default_mass()\n",
    "    g = - hopper.gravity_field().gravity_vector()[-1]\n",
    "    \n",
    "    return k, mb, ml, g"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Hopping Cycle\n",
    "\n",
    "Let's analyze the four phases of the hopping cycle illustrated below.\n",
    "We start from the **touchdown** instant $t_{\\mathrm{td}}$ (i.e. the transition from the flight phase to the stance phase).\n",
    "For the moment we assume $\\delta (t)$ to be a nonnegative constant for all $t$ (zero in the figure).\n",
    "\n",
    "At touchdown the leg has an impact with the ground and $z_{\\mathrm{l}} (t_{\\mathrm{td}}) = 0$.\n",
    "We assume this impact to be inelastic, i.e., the leg does not bounce back after the collision.\n",
    "In math: $\\dot z_{\\mathrm{l}} (t_{\\mathrm{td}}^+) = 0$, where $t_{\\mathrm{td}}^+$ denotes the instant of time immediately after touchdown.\n",
    "Because of the spring between the leg and the body, the impulse is not transmitted to the body which continues its descent with the same velocity $\\dot z_{\\mathrm{b}} (t_{\\mathrm{td}}^+) = \\dot z_{\\mathrm{b}} (t_{\\mathrm{td}}^-) < 0$.\n",
    "\n",
    "During stance, the leg is pushed towards the ground by the spring force (and gravity) hence, given that $\\dot z_{\\mathrm{l}} (t_{\\mathrm{td}}^+) = 0$, the leg-ground contact cannot be broken: $z_{\\mathrm{l}} (t) = \\dot z_{\\mathrm{l}} (t) = 0$ for all $t$ during stance.\n",
    "In contrast, the body first moves towards the ground, and then moves up towards the the hard stop of the prismatic joint.\n",
    "We call **bottom** the moment during stance when the body has minimum altitude and zero velocity, and we denote it as $t_{\\mathrm{bt}}$.\n",
    "\n",
    "When the body reaches the prismatic-joint hard stop, we have a second inelastic collision.\n",
    "We call this instant **lift off**, denoted as $t_{\\mathrm{lo}}$, since this impact between the body and the hard stop makes the leg lift off the ground.\n",
    "An instant after this inelastic collision, the body and the leg have the same velocity $\\dot z_{\\mathrm{b}} (t_{\\mathrm{lo}}^+) = \\dot z_{\\mathrm{l}} (t_{\\mathrm{lo}}^+)$.\n",
    "\n",
    "The flying phase starts after lift off.\n",
    "Our goal is to ensure that at the time $t_{\\mathrm{ap}}$, when the body is at the **apex** of its trajectory, we have $z_{\\mathrm{l}} (t_{\\mathrm{ap}}) = z_{\\mathrm{b}} (t_{\\mathrm{ap}}) = h$.\n",
    "Note that, during flight, there is no reason for the body to break contact with the hard stop and $l(t)$ is constantly zero.\n",
    "In fact, in the flight phase, the body and the leg have the same initial velocity (the one after lift off), they are both subject to the same acceleration (gravity), and the spring can only push them together (if $\\delta (t) > 0$).\n",
    "\n",
    "<img src=\"https://raw.githubusercontent.com/RussTedrake/underactuated/master/figures/exercises/one_d_hopping_cycle.jpg\" width=\"800\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Controller\n",
    "\n",
    "The inelastic impacts we have at touchdown $t_{\\mathrm{td}}$ and lift off $t_{\\mathrm{lo}}$ cause two drops in the system energy.\n",
    "The goal of the controller is to inject energy in the system to balance these losses, and ensure that the hopping height is exactly $h$.\n",
    "We achieve this by increasing the preload $\\delta (t)$ or, in other words, by increasing the potential energy stored in the spring.\n",
    "\n",
    "The plan is to modify the spring preload $\\delta(t)$ only twice per hopping cycle:\n",
    "- At bottom time $t_{\\mathrm{bt}}$ we set the preload to $\\delta (t_{\\mathrm{bt}}^+) = \\bar \\delta \\geq 0$, with $\\bar \\delta$ chosen so that the hopping height is exactly $h$.\n",
    "- At apex time $t_{\\mathrm{ap}}$ we reset the preload to $\\delta (t_{\\mathrm{ap}}^+) = 0$.\n",
    "\n",
    "To choose the correct value of $\\bar \\delta$, we will analyze the energy fluctuations during the hopping cycle.\n",
    "First, however, let's do some coding to be able to run a simulation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Spring Block in the Diagram\n",
    "\n",
    "In this notebook we \"cheat\" a bit and, instead of having a true spring between the body and the leg directly in the `urdf` file of the hopper, we mimic the elastic force using an actuator between these bodies.\n",
    "(This makes the logic behind the variable preload much easier to code.)\n",
    "This actuator is called `spring` and is one of the blocks in the diagram below.\n",
    "As inputs, it receives the `hopper` state and the preload $\\delta(t)$ from the `preload_controller`.\n",
    "As output, it sends the spring force $- k [l(t) + \\delta(t)]$ to the `hopper`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Spring(VectorSystem):\n",
    "\n",
    "    def __init__(self):\n",
    "        \n",
    "        # 5 inputs: hopper state and spring preload delta\n",
    "        # 1 output: spring force\n",
    "        VectorSystem.__init__(self, 5,  1)\n",
    "        \n",
    "        # spring stiffness\n",
    "        self.k = get_hopper_parameters()[0]\n",
    "\n",
    "    # note that this function is called at each time step\n",
    "    def DoCalcVectorOutput(self, context, state_delta, unused, spring_force):\n",
    "        \n",
    "        # set the output of the spring block to\n",
    "        # the value of the elastic force\n",
    "        l = state_delta[1]\n",
    "        delta = state_delta[4]\n",
    "        spring_force[:] = [- self.k * (l + delta)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Block Diagram and the Simulator\n",
    "\n",
    "We now write a couple of function to build the block diagram and setup a simulation environment.\n",
    "These take the argument `preload_controller` which is the controller we will synthesize later in the notebook.\n",
    "For the moment, we just need to know that `preload_controller` takes the hopper state and returns the preload $\\delta(t)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_block_diagram(preload_controller):\n",
    "\n",
    "    # initialize block diagram\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    # instantiate plant and scene graph\n",
    "    # the nonzero time_step is more robust in the handling\n",
    "    # of collisions and simplifies control synthesis (wrt time_step=0)\n",
    "    hopper, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
    "\n",
    "    # parse the hopper urdf\n",
    "    urdf_path = FindResource('models/one_d_hopper.urdf')\n",
    "    Parser(hopper).AddModelFromFile(urdf_path)\n",
    "    hopper.Finalize()\n",
    "    hopper.set_name('hopper')\n",
    "    \n",
    "    # controller for the spring preload\n",
    "    # wire hopper -> controller\n",
    "    preload_controller = builder.AddSystem(preload_controller)\n",
    "    builder.Connect(hopper.get_state_output_port(), preload_controller.get_input_port(0))\n",
    "    \n",
    "    # mux to merge the hopper state and the preload in a single cable\n",
    "    # wire hopper -> mux and controller -> mux\n",
    "    mux = builder.AddSystem(Multiplexer([4, 1]))\n",
    "    builder.Connect(hopper.get_state_output_port(), mux.get_input_port(0))\n",
    "    builder.Connect(preload_controller.get_output_port(0), mux.get_input_port(1))\n",
    "    mux.set_name('mux')\n",
    "    \n",
    "    # block that mimics the spring\n",
    "    # wire mux -> spring and spring -> hopper\n",
    "    spring = builder.AddSystem(Spring())\n",
    "    builder.Connect(mux.get_output_port(0), spring.get_input_port(0))\n",
    "    builder.Connect(spring.get_output_port(0), hopper.get_actuation_input_port())\n",
    "    spring.set_name('spring')\n",
    "    \n",
    "    # logger that records the hopper state and the preload\n",
    "    logger = builder.AddSystem(SignalLogger(5))\n",
    "    builder.Connect(mux.get_output_port(0), logger.get_input_port(0))\n",
    "    logger.set_name('logger')\n",
    "\n",
    "    # visualizer for the MultibodyPlant\n",
    "    # wire scene_graph -> visualizer\n",
    "    visualizer = builder.AddSystem(PlanarSceneGraphVisualizer(\n",
    "        scene_graph,\n",
    "        xlim=[-1, 1],\n",
    "        ylim=[-.2, 2.5],\n",
    "        show=False\n",
    "    ))\n",
    "    builder.Connect(scene_graph.get_pose_bundle_output_port(), visualizer.get_input_port(0))\n",
    "    visualizer.set_name('visualizer')\n",
    "\n",
    "    # finalize block diagram\n",
    "    diagram = builder.Build()\n",
    "    \n",
    "    return diagram"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Just to be sure that we wired things correctly, let's plot the block diagram."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dummy controller just to build the diagram for the plot\n",
    "preload_controller = MatrixGain(np.zeros((1, 4)))\n",
    "preload_controller.set_name('preload controller')\n",
    "\n",
    "# plot block diagram\n",
    "plt.figure(figsize=(20, 8))\n",
    "diagram = build_block_diagram(preload_controller)\n",
    "diagram.set_name('Block Diagram of the One-Dimensional Hopper')\n",
    "plot_system_graphviz(diagram)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The final piece of code before we can see the hopper moving is the simulator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def simulate(initial_state, duration, preload_controller):\n",
    "    \n",
    "    # get diagram\n",
    "    diagram = build_block_diagram(preload_controller)\n",
    "\n",
    "    # start recording the video for the animation of the simulation\n",
    "    visualizer = diagram.GetSubsystemByName('visualizer')\n",
    "    visualizer.start_recording()\n",
    "\n",
    "    # set up a simulator\n",
    "    simulator = Simulator(diagram)\n",
    "    simulator.set_publish_every_time_step(False)\n",
    "\n",
    "    # set the initial conditions\n",
    "    hopper = diagram.GetSubsystemByName('hopper')\n",
    "    simulator_context = simulator.get_mutable_context()\n",
    "    hopper_context = diagram.GetMutableSubsystemContext(hopper, simulator_context)\n",
    "    hopper.SetPositionsAndVelocities(hopper_context, initial_state)\n",
    "\n",
    "    # simulate from zero to duration\n",
    "    simulator.Initialize()\n",
    "    simulator.AdvanceTo(duration)\n",
    "\n",
    "    # stop the video and build the animation\n",
    "    visualizer.stop_recording()\n",
    "    animation = visualizer.get_recording_as_animation()\n",
    "\n",
    "    return animation, diagram.GetSubsystemByName('logger')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hopper Mechanical Energy\n",
    "\n",
    "We can now start the energy analysis of the hopping cycle.\n",
    "The total mechanical energy stored in the hopper at a generic time $t$ has three components (kinetic, gravitational, and elastic), and is given by\n",
    "$$\n",
    "E(t)\n",
    "=\n",
    "\\frac{1}{2} m_{\\mathrm{b}} \\dot z_{\\mathrm{b}}^2 (t)\n",
    "+\n",
    "\\frac{1}{2} m_{\\mathrm{l}} \\dot z_{\\mathrm{l}}^2 (t)\n",
    "+\n",
    "m_{\\mathrm{b}} g z_{\\mathrm{b}} (t)\n",
    "+\n",
    "m_{\\mathrm{l}} g z_{\\mathrm{l}} (t)\n",
    "+\n",
    "\\frac{1}{2} k [l (t) + \\delta(t)]^2.\n",
    "$$\n",
    "Where $l (t) = z_{\\mathrm{l}} (t) - z_{\\mathrm{b}} (t)$ is always nonnegative because of the hard limit of the prismatic joint.\n",
    "\n",
    "To support our analysis with a plot, in the next cell you are asked to write a function that computes the mechanical energy $E(t)$ given the `state` of the hopper, and the preload `delta` of the spring.\n",
    "\n",
    "Notice that:\n",
    "- We decided to represent the hopper state with $[z_{\\mathrm{b}}, l, \\dot z_{\\mathrm{b}}, \\dot l]^T$.\n",
    "- The function `mechanical_energy` should accept scalars, but also vectorized inputs coming from the logger. If we pass a $4 \\times N$ array for the `state` and an $N$ array for `delta`, with $N$ number of time steps in the simulation, we want `E` to be an $N$ array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mechanical_energy(state, delta):\n",
    "    \n",
    "    # unpack parameters and state\n",
    "    k, mb, ml, g = get_hopper_parameters()\n",
    "    zb, l, zb_dot, l_dot = state\n",
    "\n",
    "    # (vectorized) mechanical energy of the hopper\n",
    "    E = np.zeros_like(delta) # modify here\n",
    "    \n",
    "    return E"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's simulate the hopper from the initial conditions $z_{\\mathrm{b}} (0) = 1$ and $l (0) = \\dot z_{\\mathrm{b}} (0) = \\dot l (0) = 0$ for a duration of 5 seconds.\n",
    "\n",
    "Then, we plot the energy as a function of time.\n",
    "This plot is just to have a grasp on the hopper physics: for the moment we run the simulation with the spring preload $\\delta (t)$ set to zero for all times $t$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# simulation parameters\n",
    "initial_state = [1, 0, 0, 0]\n",
    "duration = 5\n",
    "\n",
    "# function to ensure that at the given state\n",
    "# there is not penetration between the rigid bodies\n",
    "def raise_if_penetration(state):\n",
    "    \n",
    "    # ground contact\n",
    "    if state[0] + state[1] < 0:\n",
    "        raise ValueError('Leg in penetration with ground.')\n",
    "        \n",
    "    # hard limit of the prismatic joint\n",
    "    if state[1] < 0:\n",
    "        raise ValueError('Body in penetration with hard joint limit.')\n",
    "        \n",
    "# check initial state\n",
    "raise_if_penetration(initial_state)\n",
    "\n",
    "# set the preload to zero independently of the hopper state\n",
    "preload_controller = MatrixGain(np.zeros((1, 4)))\n",
    "\n",
    "# run the simulation\n",
    "animation, logger = simulate(initial_state, duration, preload_controller)\n",
    "\n",
    "# evaluate mechanical energy\n",
    "time = logger.sample_times()\n",
    "state = logger.data()[:4]\n",
    "delta = logger.data()[4]\n",
    "E = mechanical_energy(state, delta)\n",
    "\n",
    "# plot mechanical energy\n",
    "plt.figure()\n",
    "plt.plot(time, E, label=r'$E(t)$')\n",
    "plt.xlim(0, duration)\n",
    "plt.xlabel('Time (s)')\n",
    "plt.ylabel('Hopper mechanical energy (J)')\n",
    "plt.grid(True)\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you did thing correctly, the energy plot should look almost like a stair.\n",
    "The energy is perfectly conserved when the hopper is in the flight phase (e.g. at time zero) and it has a drop at touchdown.\n",
    "Energy is more or less preserved during stance (in the following we will assume that it is preserved exactly), and it has a second drop at lift off.\n",
    "\n",
    "Once you're convinced about this behavior, we can start the derivation of the controller."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Derivation of the Controller\n",
    "\n",
    "A few cells above, we have said that the plan is to inject energy in the hopper during the bottom instant $t_{\\mathrm{bt}}$.\n",
    "And we noticed that we can do that by increasing the preload from zero to $\\bar \\delta$.\n",
    "Now we answer the big question mark: what's the precise value of $\\bar \\delta$ such that at the next apex $t_{\\mathrm{ap}}$ the hopper is at height $h$ from the ground?\n",
    "\n",
    "We need some math to link the system energy at the bottom and the one at the apex."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Energy Before and After the \"Bottom\"\n",
    "\n",
    "At the bottom instant $t_{\\mathrm{bt}}$ we have:\n",
    "- Leg steady in contact with the ground: $z_{\\mathrm{l}} (t_{\\mathrm{bt}}) = \\dot z_{\\mathrm{l}} (t_{\\mathrm{bt}}) = 0$.\n",
    "- Body steady in (unknown but measured) position: $z_{\\mathrm{b}} (t_{\\mathrm{bt}}) < 0$ and $\\dot z_{\\mathrm{b}} (t_{\\mathrm{bt}}) = 0$.\n",
    "- Preload moving from zero to the value to be decided: $\\delta (t_{\\mathrm{bt}}^-) = 0$ and $\\delta (t_{\\mathrm{bt}}^+)  = \\bar \\delta \\geq 0$.\n",
    "\n",
    "Just before the preload of the spring, the energy of the system is\n",
    "$$\n",
    "E(t_{\\mathrm{bt}}^-)\n",
    "=\n",
    "m_{\\mathrm{b}} g z_{\\mathrm{b}} (t_{\\mathrm{bt}})\n",
    "+\n",
    "\\frac{1}{2} k z_{\\mathrm{b}}^2 (t_{\\mathrm{bt}}),\n",
    "$$\n",
    "where we used $l(t_{\\mathrm{bt}}) = - z_{\\mathrm{b}} (t_{\\mathrm{bt}})$.\n",
    "Compressing instantaneously the spring by $\\bar \\delta$, we have an energy increase\n",
    "\\begin{align}\n",
    "E(t_{\\mathrm{bt}}^+)\n",
    "& =\n",
    "m_{\\mathrm{b}} g z_{\\mathrm{b}} (t_{\\mathrm{bt}})\n",
    "+\n",
    "\\frac{1}{2} k [\\bar \\delta - z_{\\mathrm{b}} (t_{\\mathrm{bt}})]^2\n",
    "\\\\\n",
    "& =\n",
    "m_{\\mathrm{b}} g z_{\\mathrm{b}} (t_{\\mathrm{bt}})\n",
    "+\n",
    "\\frac{1}{2} k \\bar \\delta^2\n",
    "+\n",
    "\\frac{1}{2} k z_{\\mathrm{b}}^2 (t_{\\mathrm{bt}})\n",
    "-\n",
    "k \\bar \\delta z_{\\mathrm{b}} (t_{\\mathrm{bt}}).\n",
    "\\end{align}\n",
    "We need to understand what's the right amount of preload $\\bar \\delta$ to make the system hop at height $h$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Energy Before \"Lift Off\"\n",
    "\n",
    "At the lift off time $t_{\\mathrm{lo}}$ (transition from stance to flight) we have:\n",
    "- Leg in contact with the ground, which has an instantaneous velocity increase due to the impact: $z_{\\mathrm{l}} (t_{\\mathrm{lo}}) = \\dot z_{\\mathrm{l}} (t_{\\mathrm{lo}}^-) = 0$ and $\\dot z_{\\mathrm{l}} (t_{\\mathrm{lo}}^+) > 0$.\n",
    "- Body in contact with the hard stop, which has an instantaneous velocity decrease: $z_{\\mathrm{b}} (t_{\\mathrm{lo}}) = 0$ and $0 < \\dot z_{\\mathrm{b}} (t_{\\mathrm{lo}}^+) < \\dot z_{\\mathrm{b}} (t_{\\mathrm{lo}}^-)$.\n",
    "- Preload equal to the value to be decided: $\\delta (t_{\\mathrm{lo}}) = \\bar \\delta$.\n",
    "\n",
    "Just before lift off, the energy is\n",
    "$$\n",
    "E(t_{\\mathrm{lo}}^-)\n",
    "=\n",
    "\\frac{1}{2} m_{\\mathrm{b}} \\dot z_{\\mathrm{b}}^2 (t_{\\mathrm{lo}}^-)\n",
    "+\n",
    "\\frac{1}{2} k \\bar \\delta^2.\n",
    "$$\n",
    "Since energy is (almost) preserved between $t_{\\mathrm{bt}}^+$ and $t_{\\mathrm{lo}}^-$, equating $E(t_{\\mathrm{bt}}^+)$ and $E(t_{\\mathrm{lo}}^-)$ we get the kinetic energy before lift off as a function of $\\bar \\delta$:\n",
    "$$\n",
    "\\frac{1}{2} m_{\\mathrm{b}} \\dot z_{\\mathrm{b}}^2 (t_{\\mathrm{lo}}^-)\n",
    "=\n",
    "m_{\\mathrm{b}} g z_{\\mathrm{b}} (t_{\\mathrm{bt}})\n",
    "+\n",
    "\\frac{1}{2} k z_{\\mathrm{b}}^2 (t_{\\mathrm{bt}})\n",
    "-\n",
    "k \\bar \\delta z_{\\mathrm{b}} (t_{\\mathrm{bt}}).\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### \"Lift-Off\" Impact\n",
    "\n",
    "Momentum is preserved during impacts.\n",
    "This allows to derive the velocity of the leg and the body just after lift off.\n",
    "We have\n",
    "$$\n",
    "m_{\\mathrm{b}} \\dot z_{\\mathrm{b}} (t_{\\mathrm{lo}}^-)\n",
    "=\n",
    "m_{\\mathrm{b}} \\dot z_{\\mathrm{b}} (t_{\\mathrm{lo}}^+)\n",
    "+\n",
    "m_{\\mathrm{l}} \\dot z_{\\mathrm{l}} (t_{\\mathrm{lo}}^+).\n",
    "$$\n",
    "If we model the collision as inelastic (not too far from reality and from what Drake does), the leg and the body will stick together after lift off.\n",
    "Thus, we obtain\n",
    "$$\n",
    "\\dot z_{\\mathrm{b}} (t_{\\mathrm{lo}}^+)\n",
    "=\n",
    "\\dot z_{\\mathrm{l}} (t_{\\mathrm{lo}}^+)\n",
    "=\n",
    "\\frac{m_{\\mathrm{b}}}{m_{\\mathrm{b}} + m_{\\mathrm{l}}}\n",
    "\\dot z_{\\mathrm{b}} (t_{\\mathrm{lo}}^-).\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Energy After \"Lift Off\"\n",
    "\n",
    "Just after lift off, the body sticks to the leg and the energy is\n",
    "$$\n",
    "E(t_{\\mathrm{lo}}^+)\n",
    "=\n",
    "\\frac{1}{2} (m_{\\mathrm{b}} + m_{\\mathrm{l}}) \\dot z_{\\mathrm{b}}^2 (t_{\\mathrm{lo}}^+)\n",
    "+\n",
    "\\frac{1}{2} k \\bar \\delta^2,\n",
    "$$\n",
    "Using the conservation of momentum at lift off, we obtain\n",
    "$$\n",
    "E(t_{\\mathrm{lo}}^+)\n",
    "=\n",
    "\\frac{1}{2} \n",
    "\\frac{m_{\\mathrm{b}}^2}{m_{\\mathrm{b}} + m_{\\mathrm{l}}}\n",
    "\\dot z_{\\mathrm{b}}^2 (t_{\\mathrm{lo}}^-)\n",
    "+\n",
    "\\frac{1}{2} k \\bar \\delta^2.\n",
    "$$\n",
    "Moreover, from the conservation of energy between bottom and lift off, we get\n",
    "$$\n",
    "E(t_{\\mathrm{lo}}^+)\n",
    "=\n",
    "\\frac{m_{\\mathrm{b}}}{m_{\\mathrm{b}} + m_{\\mathrm{l}}}\n",
    "\\left[\n",
    "m_{\\mathrm{b}} g z_{\\mathrm{b}} (t_{\\mathrm{bt}})\n",
    "+\n",
    "\\frac{1}{2} k z_{\\mathrm{b}}^2 (t_{\\mathrm{bt}})\n",
    "-\n",
    "k \\bar \\delta z_{\\mathrm{b}} (t_{\\mathrm{bt}})\n",
    "\\right]\n",
    "+\n",
    "\\frac{1}{2} k \\bar \\delta^2.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Energy Before the \"Apex\"\n",
    "\n",
    "At the apex $t_{\\mathrm{ap}}$ we (want to) have:\n",
    "- Leg fully extended with zero velocity: $z_{\\mathrm{l}} (t_{\\mathrm{ap}}) = h$ and $\\dot z_{\\mathrm{l}} (t_{\\mathrm{ap}}) = 0$.\n",
    "- Body in contact with the hard stop with zero velocity: $z_{\\mathrm{b}} (t_{\\mathrm{ap}}) = h$ and $\\dot z_{\\mathrm{b}} (t_{\\mathrm{ap}}) = 0$.\n",
    "- Reset the preload to zero, from the still unknown value $\\bar \\delta$: $\\delta (t_{\\mathrm{ap}}^-) = \\bar \\delta$ and $\\delta (t_{\\mathrm{ap}}^+) = 0$.\n",
    "\n",
    "The energy just before the apex is\n",
    "$$\n",
    "E(t_{\\mathrm{ap}}^-)\n",
    "=\n",
    "(m_{\\mathrm{b}} + m_{\\mathrm{l}}) g h\n",
    "+\n",
    "\\frac{1}{2} k \\bar \\delta^2.\n",
    "$$\n",
    "\n",
    "Energy is preserved in the while flying, hence $E(t_{\\mathrm{ap}}^-)$ must be equal to $E(t_{\\mathrm{lo}}^+)$.\n",
    "This gives\n",
    "$$\n",
    "(m_{\\mathrm{b}} + m_{\\mathrm{l}}) g h\n",
    "=\n",
    "\\frac{m_{\\mathrm{b}}}{m_{\\mathrm{b}} + m_{\\mathrm{l}}}\n",
    "\\left[\n",
    "m_{\\mathrm{b}} g z_{\\mathrm{b}} (t_{\\mathrm{bt}})\n",
    "+\n",
    "\\frac{1}{2} k z_{\\mathrm{b}}^2 (t_{\\mathrm{bt}})\n",
    "-\n",
    "k \\bar \\delta z_{\\mathrm{b}} (t_{\\mathrm{bt}})\n",
    "\\right].\n",
    "$$\n",
    "\n",
    "Finally! This is a linear equation that allows to express the preload $\\bar \\delta$ as a function the hopping height $h$!\n",
    "We get\n",
    "$$\n",
    "\\bar \\delta\n",
    "=\n",
    "\\frac{1}{2} z_{\\mathrm{b}} (t_{\\mathrm{bt}})\n",
    "+ \\frac{m_{\\mathrm{b}} g}{k}\n",
    "- \\frac{(m_{\\mathrm{b}} + m_{\\mathrm{l}})^2}{k m_{\\mathrm{b}} z_{\\mathrm{b}} (t_{\\mathrm{bt}})} g h.\n",
    "$$\n",
    "\n",
    "Note that $\\bar \\delta$ could in principle result to be negative, when we have more energy than we need for the hop.\n",
    "In that case, we just set $\\bar \\delta = 0$, and we let the impacts in the cycle pull energy out of the system until we arrive at the desired level."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### What Happens After the \"Apex\"?\n",
    "\n",
    "Just after $t_{\\mathrm{ap}}$, we reset the spring preload $\\delta(t)$ to zero.\n",
    "You can convince yourself that this does not have any impact in the motion of th hopper.\n",
    "Before the apex the spring force is contrasted by the hard stop, the reset just zeroes this internal force.\n",
    "(Consider drawing the free body diagram for the leg and the body separately if you find this hard to believe.)\n",
    "\n",
    "The impact at touchdown can be analyzed similarly to the one at lift off (it is actually simpler) but, since it does not affect our control decision, we do not care about the details of what happens there."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulate the Closed-Loop System\n",
    "\n",
    "Finally we are ready to implement our controller!\n",
    "\n",
    "In the next cell you are asked to implement the controller we just derived by working on the method `DoCalcVectorOutput` of the class `PreloadController`.\n",
    "At each time step $t$, `DoCalcVectorOutput` must set the entry of its argument `delta` to the desired preload value $\\delta(t)$.\n",
    "To compute this value, you are allowed to use:\n",
    "- `self.h`: the desired hopping height.\n",
    "- `self.hopper_parameters`: the physical parameters of the hopper ($k$, $m_{\\mathrm{b}}$, $m_{\\mathrm{l}}$, and $g$).\n",
    "- `state`: the hopper state $[z_{\\mathrm{b}}, l, \\dot z_{\\mathrm{b}}, \\dot l]^T$ at the current time step;\n",
    "- `self.last_delta`: the preload $\\delta(t)$ you set the last time `DoCalcVectorOutput` has been called (i.e. last time step).\n",
    "- `self.last_state`: the hopper state at the last time step.\n",
    "\n",
    "The implementation of this function is nontrivial.\n",
    "Feel free to take the path you prefer to do that but, in case you feel lost, this is a sketch of our implementation.\n",
    "\n",
    "We added two methods to the class `PreloadController`:\n",
    "- `is_bottom_or_apex`. Using the `state` and the `last_state`, this method identifies if the current time step is the bottom $t_{\\mathrm{bt}}$, the apex $t_{\\mathrm{ap}}$, or neither of the two. This can be done, e.g., by checking if the velocity of the body $\\dot z_{\\mathrm{b}}$ changed sign from the last time step. Whereas, to distinguish between the bottom and the apex, you can look at the sign of the position of the body $z_{\\mathrm{b}}$.\n",
    "- `get_delta_bar`. Using `h`, the `hopper_parameters`, and the current `state`, this method computes the value of $\\bar \\delta$ using the formula derived above.\n",
    "\n",
    "Finally, `DoCalcVectorOutput` calls first `is_bottom_or_apex`.\n",
    "If the current time is the apex $t_{\\mathrm{ap}}$, it sets the entry in `delta` to zero.\n",
    "If the current time is the bottom $t_{\\mathrm{bt}}$, it calls `get_delta_bar` and it sets `delta` to $\\max(\\bar \\delta, 0)$ (remember that we want $\\delta(t) \\geq 0$).\n",
    "If none of the previous, it sets `delta` to the `last_delta`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PreloadController(VectorSystem):\n",
    "    \n",
    "    def __init__(self, h):\n",
    "        \n",
    "        # 4 inputs: hopper state\n",
    "        # 1 output: spring preload\n",
    "        VectorSystem.__init__(self, 4,  1)\n",
    "        \n",
    "        # desired hopping height\n",
    "        self.h = h\n",
    "        \n",
    "        # store physical parameters of the hopper\n",
    "        self.hopper_parameters = get_hopper_parameters()\n",
    "        \n",
    "        # the controller keeps track of the last value of\n",
    "        # the preload and the last state of the hopper\n",
    "        self.last_delta = 0\n",
    "        self.last_state = np.full(4, np.nan)\n",
    "    \n",
    "    # note that this function is called at each time step\n",
    "    def DoCalcVectorOutput(self, context, state, unused, delta):\n",
    "        \n",
    "        # troubleshooting: inside this class you must use self.h\n",
    "        # (class attribute) and not h (global variable)\n",
    "        h = self.h\n",
    "        \n",
    "        # assign value of the preload at each time step\n",
    "        delta[:] = [0] # modify here\n",
    "        \n",
    "        # update controller internal memory\n",
    "        self.last_delta = delta[0]\n",
    "        self.last_state = state\n",
    "        \n",
    "    # feel free to add methods to this class if you need to\n",
    "    # modify here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here you can test your work by playing with the initial conditions and the value of $h$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# simulation parameters\n",
    "initial_state = [.1, 0, 0, 0]\n",
    "raise_if_penetration(initial_state)\n",
    "duration = 5\n",
    "\n",
    "# desired hopping height\n",
    "h = 1\n",
    "\n",
    "# synthesize the controller\n",
    "preload_controller = PreloadController(h)\n",
    "\n",
    "# run simulation\n",
    "animation, logger = simulate(initial_state, duration, preload_controller)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Animate and Plot the Simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# play the video\n",
    "HTML(animation.to_jshtml())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# unpack logger data\n",
    "time = logger.sample_times()\n",
    "state = logger.data()[:4]\n",
    "delta = logger.data()[4]\n",
    "\n",
    "# plot body position zb(t) and target value h\n",
    "plt.figure()\n",
    "plt.plot(time, state[0], label=r'$z_{\\mathrm{b}}(t)$')\n",
    "plt.plot(time, [h] * len(time), label=r'$h$')\n",
    "plt.xlim(0, duration)\n",
    "plt.xlabel('Time (s)')\n",
    "plt.ylabel('Hopping height (m)')\n",
    "plt.grid(True)\n",
    "plt.legend()\n",
    "\n",
    "# plot spring preload delta(t)\n",
    "plt.figure()\n",
    "plt.plot(time, delta, label=r'$\\delta(t)$')\n",
    "plt.xlim(0, duration)\n",
    "plt.xlabel('Time (s)')\n",
    "plt.ylabel('Spring preload (m)')\n",
    "plt.grid(True)\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Autograding\n",
    "You can check your work by running the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from underactuated.exercises.simple_legs.one_d_hopper.test_one_d_hopper import TestOneDHopper\n",
    "from underactuated.exercises.grader import Grader\n",
    "Grader.grade_output([TestOneDHopper], [locals()], 'results.json')\n",
    "Grader.print_test_results('results.json')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}